# author: zhaofeng-shu33
import os
import argparse
import pdb
import re
import time

from PIL import Image # pillow
from docx import Document
from docx.shared import Pt
import docx
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.shared import Inches
from copy import deepcopy

# initialization of databse related variable
IMAGE_LIST = []
base_dir_name = 'output'
dir_names = ['荷园2号楼','荷园4号楼','荷园5号楼','荷园1号楼'] # omit #3
cnt_num = 0
lack_dic = {}
# read of docx files
def get_docx(dir):
    for i in os.listdir(os.path.join(base_dir_name, dir)):
        if i.find('docx') > 0:
            if(os.path.getsize(os.path.join(dir,i))>1024*502):#>0.5MB
                yield i
def get_basic_info(dir_name, file_name):
    house_number_str = re.search('[1-5]', dir_name).group(0)
    re_obj = re.search('0([2-3])[0-9]{2}-([0-9]{2})([0-9]{2})', file_name)       
    start_floor_number = int(re_obj.group(1))
    end_floor_number = int(re_obj.group(2))
    floor_index_str = re_obj.group(3)
    return (house_number_str, start_floor_number, end_floor_number, floor_index_str)

def add_formated_str(para, added_str):
    run = para.add_run(added_str)
    run.font.bold = True
    if(re.match('[0-9]+',added_str)): # suppose all are numbers
        run.font.name = 'Calibri'
    else: # suppose all are Chinese characters
        run.font.name = 'SimSun'
    run.font.size = Pt(14)

def add_info(doc, dic, floor_cnt):
    if(dic.get(floor_cnt)):
        for index,i in enumerate(dic[floor_cnt]):
            t = doc.tables[-(3-index)] # get the last table
            t.cell(0,2).text = i[0] # name
            t.cell(0,4).text = gender_dic[i[1]] # gender
            t.cell(1,2).text = i[2] # student_id
def select(table):
    student_id = None
    if(table.cell(1,2).text):
        student_id = table.cell(1,2).text.strip()
    return student_id
def update_table(table, info_item):
    name = info_item[0][0]    
    table.cell(0,2).text = name # name
    try:
        table.cell(0,4).text = gender_dic[info_item[0][1]]
    except KeyError as e:
        pdb.set_trace()
    table.cell(1,2).text = info_item[0][2].lstrip('1')
    mobile = info_item[0][3]
    if(mobile):
        table.cell(1,4).text = mobile
    major = info_item[0][4]
    if(major):
        table.cell(2,2).text = major
    group_leader = info_item[0][6]
    if(group_leader):
        table.cell(2,4).text = group_leader
    mentor_name = info_item[0][7]
    if(mentor_name):
        table.cell(3,2).text = mentor_name
    mentor_phone = info_item[1]
    if(mentor_phone):
        table.cell(3,4).text = mentor_phone
    relative_name = info_item[0][8]
    if(relative_name):
        if(re.match('[0-9]+',relative_name)):
            print('relative_name cannot contain numbers')
            pdb.set_trace()
        table.cell(4,2).text = relative_name
    relative_phone = info_item[0][9]
    if(relative_phone):
        table.cell(4,4).text = relative_phone
    house_address = info_item[0][10]
    if(house_address):
        table.cell(5,2).text = house_address
    global added_num
    added_num += 1

def iter_rId_parts(d):
    """
    Generate exactly one reference to each of the parts in the package by
    performing a depth-first traversal of the rels graph.
    """
    def walk_parts(source, visited=list()):
        for rId, rel in source.rels.items():
            if rel.is_external:
                continue
            part = rel.target_part
            if part in visited:
                continue
            visited.append(part)
            yield (rId, part)
            new_source = part
            for rId, part in walk_parts(new_source, visited):
                yield (rId, part)

    for item in walk_parts(d.part.package):
        yield item    

def remove_images(rid, doc):    
    for key, part in iter_rId_parts(doc):
        if(key == rid):
            if(type(part) is not docx.parts.image.ImagePart):
                print('error not imagepart')
            part._blob = b''
            return
    #return_value = walk_rId_parts_remove_images(rid, d.part.package)
    pdb.set_trace()
def remove_leaved(table, student_name, student_id, student_dormitory, doc):
    cursor.execute('select student_dormitory, student_id from qinghua where student_name = "%s"'%student_name)
    HAS_DORMITORY = False
    for i in cursor.fetchall():
        if i[0] is not None and i[1].find(student_id)>0:
            HAS_DORMITORY = True
            if(student_dormitory != i[0]):
                # print("Inconsistent dormitory found for {0}, {1}, {2} v.s. {3} assume no dormitory!".format(student_name, student_id, student_dormitory, i[0]))
                HAS_DORMITORY = False
            break
    if HAS_DORMITORY == True:
        return False
    para = table.cell(0,0).paragraphs[0]
    Ptr = para._p.find('.//v:imagedata',namespaces=para._p.nsmap)
    if(Ptr is not None):
        rid = Ptr.get('{%s}id'%Ptr.nsmap['r'])                                
        remove_images(rid, doc)
    else: # existing person does not have image
        pass
    para = para.clear()
    table.cell(5,2)._element.text = ''
    for i in range(0,5):
        for j in range(1,3):
            table.cell(i,2*j).text = ''
    global cnt_num    
    cnt_num += 1
    return True
def get_image(student_id):
    '''
        get the image path for the given student_id
    '''
    try:
        index = IMAGE_LIST.index(str(student_id)+'.bmp')
    except ValueError as e:
        try:
            index = IMAGE_LIST.index(str(student_id)+'.jpg')
        except ValueError as e:
            return None
    image_path = os.path.join('verification_example','img_data',IMAGE_LIST[index])
    IMAGE_LIST.pop(index)
    return image_path
def update_tables(tables, info_items, doc): #info_items = [(infos),mentor_mobile]
    if(len(info_items)>3):
        print("A dormitory cannot have more than three persons, at %s"%info_items[0][0][-1])
        return
    if(len(info_items)==0):
        return
    table_name_list = []
    for table in tables:
        table_student_name = table.cell(0,2).text.strip()
        student_id_value =  table.cell(1,2).text
        try:
            student_dom = info_items[0][0][-1]
        except Exception as e:
            pdb.set_trace()
        is_removed = remove_leaved(table, table_student_name, student_id_value, student_dom, doc)
        if(table_student_name and is_removed == False):
            table_name_list.append(table_student_name)
    for info_item in info_items:
        student_name = info_item[0][0]
        if(table_name_list.count(student_name)>0):
            continue
        table = get_available_table(tables)
        if(table is None):
            pdb.set_trace()
        update_table(table, info_item)    
    return

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--year', default='2019')
    args = parser.parse_args()

    re_str = re.compile(args.year)
    for i in os.listdir(os.path.join('verification_example','img_data')):
        if(re_str.match(i)): # only use the pictures start with 2018
            IMAGE_LIST.append(i)

    # do the work
    start_time = time.time()
    for dir_name in dir_names:
        for file_name in get_docx(dir_name):
            f = open(os.path.join(base_dir_name, dir_name, file_name), 'rb')
            d = Document(f)
            house_number_str, start_floor_number, end_floor_number, floor_index_str = get_basic_info(dir_name, file_name)
            try:
                assert(len(d.tables)==3*(end_floor_number - start_floor_number + 1))
            except AssertionError as e:
                pdb.set_trace()
            for table_index in range(len(d.tables)):
                student_id = select(d.tables[table_index]) # get the student identity number
                if(student_id):
                    image_path = get_image(student_id)
                    if(image_path):
                        exception_occur = False
                        try:
                            d.tables[table_index].cell(0,0).paragraphs[0].add_run().add_picture(image_path, Inches(2))
                        except Exception as e:
                            exception_occur = True
                        if exception_occur:
                            im = Image.open(image_path)
                            im.save(image_path)
                            try:
                                d.tables[table_index].cell(0,0).paragraphs[0].add_run().add_picture(image_path, Inches(2))
                            except Exception as e:
                                pdb.set_trace()
                        cnt_num += 1
                
            d.save(os.path.join(base_dir_name, dir_name, file_name))
    end_time = time.time()
    print("Statistics: ")
    print("%d pictures added"%cnt_num)

    print("total time used: %0.2f"%(end_time - start_time))